package com.example.userdemo.util;

import com.opencsv.CSVReader;
import com.opencsv.CSVWriter;

import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.List;

/**
 * @Description
 * @Author liqinglin
 * @Date 2023/12/8 15:29
 */
public class CSVUtil {
    public static void group(){
        try {
            // 读取CSV文件
            CSVReader reader = new CSVReader(new FileReader("D:\\project\\上线准备\\emllm\\12-10刷数据\\预训练列表.csv"));
            String[] nextLine;
            List<String> taskIds = new ArrayList<>();
            List<Integer> counts = new ArrayList<>();

            reader.readNext();

            // 读取task_id和count字段
            while ((nextLine = reader.readNext()) != null) {
                taskIds.add(nextLine[0]);
                counts.add(Integer.parseInt(nextLine[1]));
            }

            // 输出每组的第一个task_id和最后一个task_id以及累计和
            // 写入结果到CSV文件
            CSVWriter writer = new CSVWriter(new FileWriter("output.csv"));
            int sum = 0;
            String firstTaskId = "";
            String lastTaskId = "";
            for (int i = 0; i < taskIds.size(); i++) {
                if (sum + counts.get(i) <= 50000) {
                    if (sum == 0) {
                        firstTaskId = taskIds.get(i);
                    }
                    sum += counts.get(i);
                    lastTaskId = taskIds.get(i);
                } else {
                    String[] data = {firstTaskId, lastTaskId, String.valueOf(sum)};
                    writer.writeNext(data);
                    sum = counts.get(i);
                    firstTaskId = taskIds.get(i);
                }
            }
            String[] data = {firstTaskId, lastTaskId, String.valueOf(sum)};
            writer.writeNext(data);
            writer.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
